from backbones.simple_rnn import RNN
from dataset.linear_loader import generate_sin_loader
from torch import nn
import torch

if __name__ == '__main__':
    num = 100  # 数据样本的总数量
    batch_size = 10  # 数据批次大小
    seq_len = 5  # 输入输出序列长度
    skip = 1  # 输出输出序列的间隔
    dataloader = generate_sin_loader(num=num, batch_size=batch_size, seq_len=5, skip=1)
    vocab_size = 1
    hidden_size = 20
    device = torch.device("cuda")
    model = RNN(vocab_size, hidden_size, device=device)
    criterion = nn.MSELoss()
    # model.parameters()在nn.Linear与nn.Conv2d都是默认含有w,b参数，
    # 由于自定义的RNN中参数是自定义的，因此对于模型而言该参数并未注册（nn.Parameter）。
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

    epochs = 200
    for epoch in range(epochs):
        loss_list = []
        for inputs, outputs in dataloader:
            optimizer.zero_grad()
            # inputs和outputs的实际形状 (batch_size, seq_len(时间序列), vocab_size=1)
            inputs = inputs.unsqueeze(-1).transpose(0, 1)  # 将时间序列的维度提前
            outputs = outputs.unsqueeze(-1).transpose(0, 1)  # 将时间序列的维度提前
            inputs = inputs.to(device)  # 切换数据设备
            outputs = outputs.to(device)  # 切换数据设备
            # 由于每个批次的数据可能是不相关的，因此隐藏状态在每一次的重新运行中都需要初始化
            h0 = torch.randn((batch_size, hidden_size), device=device)  # 初始化的隐藏状态
            predicts, _ = model(inputs, h0)

            loss = criterion(predicts, outputs)
            loss.backward()

            # 梯度裁剪：因为RNN没有CNN的BatchNorm
            # 将梯度值大于0.5的部分全部设置为0
            nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)

            optimizer.step()
            loss_list.append(loss.item())

        avg_loss = sum(loss_list) / len(loss_list)
        print(f"epoch:{epoch + 1}/{epochs} -- loss:{avg_loss:.4f}")
